import os
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
 
from swift.llm import DatasetName, ModelType, SftArguments, sft_main
 
sft_args = SftArguments(
    model_type=ModelType.qwen-7b-chat,
    dataset=[DatasetName.blossom-math-zh],
    train_dataset_sample=1000,
    logging_steps=5,
    max_length=2048,
    learning_rate=5e-5,
    warmup_ratio=0.4,
    output_dir='output',
    lora_target_modules=['ALL'],
output = sft_main(sft_args)
best_model_checkpoint = output['best_model_checkpoint']
print(f'best_model_checkpoint: {best_model_checkpoint}')